# util/pdf_util.py - PDF处理工具类

import fitz
import pandas as pd
import os
from typing import List, Dict, Any
import logging


logger = logging.getLogger("doc_util")


def load_pdf_documents(file_path):
    """
    专门加载PDF文档的函数 - 增强版（使用表格替换方法）
    
    Args:
        file_path: PDF文件路径
        
    Returns:
        list: Document对象列表
    """
    try:
        from llama_index.core import Document
        
        print(f"\n📄 处理PDF文件: {file_path}")
        print("=" * 60)
        
        # 使用表格提取方法
        docs = _simple_table_extraction(file_path)
        
        if not docs:
            logger.warning(f"未从PDF文件中提取到内容: {file_path}")
            return []
        
        # 获取表格替换后的完整文档内容
        document_text = _get_document_with_tables_replaced(docs)
        
        if not document_text.strip():
            logger.warning(f"PDF文件内容为空: {file_path}")
            return []
        
        # 创建Document对象
        total_tables = sum(doc['table_count'] for doc in docs)
        
        document_obj = Document(
            text=document_text,
            metadata={
                "file_path": file_path,
                "page_count": len(docs),
                "file_type": "pdf",
                "has_tables": total_tables > 0,
                "table_count": total_tables,
                "processing_method": "table_replacement"
            }
        )
        
        print(f"✅ PDF处理完成: 共{len(docs)}页，总表格数: {total_tables}")
        print(f"📝 总文本长度: {len(document_text)} 字符")
        print("=" * 60)
        
        return [document_obj]
        
    except ImportError:
        logger.warning("PyMuPDF未安装，使用回退方案处理PDF")
        # 回退到PyPDF2
        try:
            import PyPDF2
            from llama_index.core import Document
            
            documents = []
            with open(file_path, 'rb') as file:
                pdf_reader = PyPDF2.PdfReader(file)
                
                for page_num, page in enumerate(pdf_reader.pages):
                    text = page.extract_text()
                    if text.strip():
                        doc = Document(
                            text=text,
                            metadata={
                                "file_path": file_path,
                                "page_number": page_num + 1,
                                "file_type": "pdf",
                                "has_tables": False,  # PyPDF2无法检测表格
                                "table_count": 0
                            }
                        )
                        documents.append(doc)
            
            return documents
            
        except ImportError:
            logger.error("PyPDF2未安装，请运行: pip install PyPDF2")
            # 最终回退到SimpleDirectoryReader
            from llama_index.core import SimpleDirectoryReader
            return SimpleDirectoryReader(input_files=[file_path]).load_data()
    
    except Exception as e:
        logger.error(f"PDF处理失败 {file_path}: {e}")
        # 回退到SimpleDirectoryReader
        from llama_index.core import SimpleDirectoryReader
        return SimpleDirectoryReader(input_files=[file_path]).load_data()
    
def load_jsonl_documents(file_path):
    """
    专门加载JSONL文档的函数
    
    Args:
        file_path: JSONL文件路径
        
    Returns:
        list: Document对象列表
    """
    import json
    from llama_index.core import Document
    
    documents = []
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            for line_num, line in enumerate(f, 1):
                line = line.strip()
                if not line:
                    continue
                
                try:
                    data = json.loads(line)
                    # 假设JSONL中每个对象有'text'字段
                    if 'text' in data and data['text'].strip():
                        doc = Document(
                            text=data['text'],
                            metadata={
                                "file_path": file_path,
                                "line_number": line_num,
                                "file_type": "jsonl",
                                **data.get('metadata', {})
                            }
                        )
                        documents.append(doc)
                    else:
                        logger.warning(f"JSONL行 {line_num} 缺少text字段或为空")
                        
                except json.JSONDecodeError as e:
                    logger.warning(f"JSONL行 {line_num} 解析失败: {e}")
    
    except Exception as e:
        logger.error(f"JSONL文件处理失败 {file_path}: {e}")
    
    return documents

def _convert_to_markdown_table(df):
    """
    转换为Markdown表格格式（LLM容易理解）
    """
    if df.empty:
        return ""
    
    # 表头
    headers = list(df.columns)
    md_table = "| " + " | ".join(headers) + " |\n"
    
    # 分隔线
    md_table += "| " + " | ".join(["---"] * len(headers)) + " |\n"
    
    # 数据行
    for _, row in df.iterrows():
        row_data = [str(row[col]) for col in df.columns]
        md_table += "| " + " | ".join(row_data) + " |\n"
    
    return md_table

def _convert_to_natural_language_generic(df):
    """
    通用自然语言转换
    """
    if df.empty:
        return "空表格"
    
    natural_text = f"这是一个{len(df)}行{len(df.columns)}列的表格：\n\n"
    
    # 描述表头
    headers = list(df.columns)
    natural_text += f"表格包含以下字段：{', '.join(headers)}\n\n"
    
    # 逐行描述数据
    natural_text += "数据内容：\n"
    for idx, row in df.iterrows():
        row_items = []
        for col in df.columns:
            if str(row[col]).strip():  # 只包含非空值
                row_items.append(f"{col}：{row[col]}")
        
        if row_items:
            natural_text += f"- 第{idx+1}行: {'; '.join(row_items)}\n"
    
    return natural_text

def _simple_table_extraction(pdf_path):
    """
    简化版表格提取，专注于基本功能
    返回包含文本和表格（Markdown格式）的文档列表
    """
    doc = fitz.open(pdf_path)
    all_docs = []
    
    for page_num in range(len(doc)):
        page = doc.load_page(page_num)
        
        print(f"\n📄 第 {page_num + 1} 页:")
        print("-" * 40)
        
        # 获取页面完整文本
        full_text = page.get_text()
        
        # 直接使用PyMuPDF的表格检测
        tables = page.find_tables()
        
        if tables.tables:
            print(f"检测到 {len(tables.tables)} 个表格")
            
            # 处理表格
            table_replacements = []
            
            for i, table in enumerate(tables.tables):
                try:
                    df = table.to_pandas()
                    print(f"表格 {i+1}: {len(df)}行 x {len(df.columns)}列")
                    
                    # 转换为Markdown格式
                    md_table = _convert_to_markdown_table(df)
                    
                    table_replacements.append({
                        "markdown": md_table,
                        "natural_language": _convert_to_natural_language_generic(df)
                    })
                    
                except Exception as e:
                    print(f"处理表格 {i+1} 时出错: {e}")
            
            all_docs.append({
                "page": page_num + 1,
                "content": full_text,  # 保持原始文本
                "table_count": len(tables.tables),
                "tables": table_replacements
            })
            
        else:
            print("未检测到表格")
            
            # 如果没有表格，直接使用原始文本
            all_docs.append({
                "page": page_num + 1,
                "content": full_text,
                "table_count": 0,
                "tables": []
            })
            
            # 显示页面文本内容供参考
            if full_text.strip():
                print("页面文本内容预览:")
                print(full_text[:1000] + "..." if len(full_text) > 1000 else full_text)
    
    doc.close()
    return all_docs

def _get_document_with_tables_replaced(docs):
    """
    整理文档内容，将检测到的表格直接替换为Markdown格式
    返回整理好的文本内容，不直接保存文件
    
    Args:
        docs: _simple_table_extraction返回的文档列表
        
    Returns:
        str: 包含表格Markdown格式的完整文档文本
    """
    document_text = ""
    
    for doc in docs:
        if doc['table_count'] > 0:
            # 对于有表格的页面，使用智能表格替换
            content = doc['content']
            processed_content = _smart_table_replacement(content, doc['tables'])
            
            # 处理后的内容，过滤掉多余空行
            compact_content = '\n'.join([line for line in processed_content.split('\n') if line.strip()])
            document_text += compact_content
        else:
            # 对于无表格页面，直接使用内容并过滤多余空行
            content = doc['content']
            compact_content = '\n'.join([line for line in content.split('\n') if line.strip()])
            document_text += compact_content
        
        # 页与页之间只加一个换行符
        document_text += "\n"
    
    return document_text

def _smart_table_replacement(content, tables):
    """
    智能表格替换：基于文本结构和表格内容匹配
    """
    lines = content.split('\n')
    
    # 如果只有一个表格，使用简化逻辑
    if len(tables) == 1:
        return _replace_single_table(lines, tables[0])
    
    # 多个表格的情况
    return _replace_multiple_tables(lines, tables)

def _replace_single_table(lines, table):
    """
    替换单个表格：找到最可能包含表格内容的区域
    """
    # 分析文本结构，找到最密集的数据区域
    data_regions = _find_data_dense_regions(lines)
    
    if data_regions:
        # 使用最大的数据区域
        start, end = data_regions[0]
        return _replace_region_with_table(lines, start, end, table)
    else:
        # 如果没有明显的数据区域，在文档中段插入表格
        return _insert_table_at_optimal_position(lines, table)

def _replace_multiple_tables(lines, tables):
    """
    替换多个表格：基于内容相似度匹配
    """
    # 找到所有可能的数据密集区域
    data_regions = _find_data_dense_regions(lines)
    
    # 如果区域数量与表格数量匹配，直接对应替换
    if len(data_regions) == len(tables):
        result_lines = []
        last_end = 0
        
        for (start, end), table in zip(data_regions, tables):
            # 添加区域前的文本
            result_lines.extend(lines[last_end:start])
            # 添加表格
            result_lines.append(table['markdown'])
            last_end = end + 1
        
        # 添加剩余文本
        result_lines.extend(lines[last_end:])
        return '\n'.join(result_lines)
    
    # 否则使用基于内容的匹配
    return _replace_tables_by_content_similarity(lines, tables)

def _find_data_dense_regions(lines):
    """
    找到数据密集的区域（可能包含表格）
    基于以下通用特征：
    1. 行长度的一致性
    2. 词汇密度的变化
    3. 空白字符的分布
    """
    regions = []
    current_region = None
    
    for i, line in enumerate(lines):
        line = line.strip()
        
        if _is_likely_table_line(line, _get_context_lines(lines, i)):
            if current_region is None:
                current_region = [i, i]
            else:
                current_region[1] = i
        else:
            if current_region is not None:
                # 确保区域有足够的大小
                region_size = current_region[1] - current_region[0] + 1
                if region_size >= 2:  # 至少2行
                    regions.append(tuple(current_region))
                current_region = None
    
    # 处理最后一个区域
    if current_region is not None:
        region_size = current_region[1] - current_region[0] + 1
        if region_size >= 2:
            regions.append(tuple(current_region))
    
    # 按区域大小排序
    regions.sort(key=lambda x: x[1] - x[0], reverse=True)
    return regions

def _is_likely_table_line(line, context_lines):
    """
    判断一行是否可能是表格行
    """
    if not line or len(line) < 2:
        return False
    
    # 特征1: 行结构一致性（与上下文比较）
    if _has_consistent_structure(line, context_lines):
        return True
    
    # 特征2: 包含多个分隔的字段
    field_count = _estimate_field_count(line)
    if field_count >= 2 and field_count <= 8:  # 合理的字段数量范围
        return True
    
    # 特征3: 数据密度（数字、短词的比例）
    if _has_high_data_density(line):
        return True
    
    return False

def _has_consistent_structure(line, context_lines):
    """
    检查与上下文行的结构是否一致
    """
    if not context_lines:
        return False
    
    current_field_count = estimate_field_count(line)
    
    # 检查上下文中是否有相似结构的行
    similar_structure_count = 0
    for context_line in context_lines:
        context_field_count = estimate_field_count(context_line)
        if abs(current_field_count - context_field_count) <= 1:  # 字段数量相近
            similar_structure_count += 1
    
    # 如果多数上下文行结构相似，可能是表格
    return similar_structure_count >= len(context_lines) * 0.6

def _estimate_field_count(line):
    """
    估计一行中的字段数量
    """
    # 方法1: 基于分隔符
    if '\t' in line:
        return len([f for f in line.split('\t') if f.strip()])
    
    # 方法2: 基于多个连续空格
    if '  ' in line:
        parts = line.split('  ')
        return len([p for p in parts if p.strip()])
    
    # 方法3: 基于竖线（Markdown表格）
    if '|' in line and line.count('|') >= 2:
        return line.count('|') - 1
    
    # 方法4: 基于单词数量（保守估计）
    words = line.split()
    return min(len(words), 10)  # 限制最大字段数

def _has_high_data_density(line):
    """
    检查行是否有高数据密度
    """
    if not line:
        return False
    
    # 计算数字字符比例
    digit_chars = sum(1 for c in line if c.isdigit())
    # 计算短词比例（表格中通常有较多短词）
    words = line.split()
    short_words = sum(1 for word in words if 1 <= len(word) <= 8)
    
    total_chars = len(line)
    total_words = len(words)
    
    if total_chars > 0 and total_words > 0:
        digit_ratio = digit_chars / total_chars
        short_word_ratio = short_words / total_words
        
        # 高数据密度特征：有一定比例的数字和较多短词
        return digit_ratio > 0.1 or short_word_ratio > 0.7
    
    return False

def _get_context_lines(lines, current_index, window_size=2):
    """
    获取当前行的上下文行
    """
    start = max(0, current_index - window_size)
    end = min(len(lines), current_index + window_size + 1)
    return [lines[i].strip() for i in range(start, end) if i != current_index]

def _replace_region_with_table(lines, start, end, table):
    """
    用表格替换指定区域
    """
    result_lines = []
    # 添加区域前的文本
    result_lines.extend(lines[:start])
    # 添加表格
    result_lines.append(table['markdown'])
    # 添加区域后的文本
    result_lines.extend(lines[end + 1:])
    return '\n'.join(result_lines)

def _insert_table_at_optimal_position(lines, table):
    """
    在最佳位置插入表格
    """
    # 寻找段落边界
    paragraph_boundaries = _find_paragraph_boundaries(lines)
    
    if paragraph_boundaries:
        # 在第一个段落边界后插入
        insert_pos = paragraph_boundaries[0]
        result_lines = lines[:insert_pos + 1] + [table['markdown']] + lines[insert_pos + 1:]
    else:
        # 在文档中段插入
        mid_point = len(lines) // 2
        result_lines = lines[:mid_point] + [table['markdown']] + lines[mid_point:]
    
    return '\n'.join(result_lines)

def _find_paragraph_boundaries(lines):
    """
    找到段落边界（空行）
    """
    boundaries = []
    for i in range(1, len(lines) - 1):
        if not lines[i].strip() and lines[i-1].strip() and lines[i+1].strip():
            boundaries.append(i)
    return boundaries

def _replace_tables_by_content_similarity(lines, tables):
    """
    基于内容相似度替换多个表格
    """
    # 简化的实现：在发现的数据区域插入表格
    data_regions = _find_data_dense_regions(lines)
    result_lines = lines.copy()
    
    # 按区域大小降序排列，优先替换大的区域
    data_regions.sort(key=lambda x: x[1] - x[0], reverse=True)
    
    tables_used = 0
    for start, end in data_regions:
        if tables_used < len(tables):
            # 替换这个区域
            result_lines = (
                result_lines[:start] + 
                [tables[tables_used]['markdown']] + 
                result_lines[end + 1:]
            )
            tables_used += 1
    
    # 如果还有表格未使用，在文档末尾添加
    for i in range(tables_used, len(tables)):
        result_lines.append(tables[i]['markdown'])
    
    return '\n'.join(result_lines)

